from utils.node import Node
from utils.global_variables import *


def is_matrix(data: Data, mat):
    if mat in data.global_vars.keys():
        return True
    elif mat in data.tmp_dict.keys():
        True
    elif mat in data.const_dict.keys():
        return True
    else:
        return False


def get_matrix_pointer(data: Data, mat: str) -> str:
    # 需要保证data里必须有mat

    var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
    dim = var_dicts[mat]['dim']
    if dim == '1*1':
        return f'&{mat}'
    dims = dim.split('*')
    if len(dims) == 2 and '1' in dims:
        return f'&{mat}[0]'
    return f'&{mat}' + '[0]' * len(dims)


def matrix_multiply_dims_match(dim1, dim2, mat1, mat2):
    dim1, dim2 = dim1.split('*'), dim2.split('*')
    if len(dim1) != 2 or len(dim2) != 2:
        erro_print(traceback.extract_stack(),
                   'Matrix dim is wrong!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None
    # if dim1[-1] == '1':
    #     dim1 = ['1', dim1[0]]
    # if dim2[0] == '1':
    #     dim2 = [dim2[-1], '1']
    if dim1[-1] != dim2[0]:
        erro_print(traceback.extract_stack(),
                   'Matrices not matched for matrix multiply!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None
    res_dim = '{}*{}*{}'.format(dim1[0], dim1[-1], dim2[-1])
    return res_dim


def matrix_multiScalar_dim(dim, mat):
    dim = dim.split('*')
    if len(dim) != 2:
        erro_print(traceback.extract_stack(), 'Matrix dim is wrong!({}:{})'.format(mat, str(dim)))
        return None
    return dim


def matrix_add_sub_dims_match(dim1, dim2, mat1, mat2):
    dim1, dim2 = dim1.split('*'), dim2.split('*')
    if len(dim1) != 2 or len(dim2) != 2:
        erro_print(traceback.extract_stack(),
                   'Matrix dim is wrong!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None
    if '1' in dim1 and '1' in dim2:
        if dim1[0] == '1' and dim2[0] != '1':
            dim1 = [dim1[1], dim1[0]]
        elif dim1[1] == '1' and dim2[1] != '1':
            dim2 = [dim2[1], dim2[0]]
    if dim1[0] != dim2[0] or dim1[-1] != dim2[-1]:
        erro_print(traceback.extract_stack(),
                   'Matrices not matched for matrix add or sub!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None
    res_dim = '{}*{}'.format(dim1[0], dim1[-1])
    return res_dim


def matrix_dims_match(dim1, dim2, mat1, mat2):
    dim1, dim2 = dim1.split('*'), dim2.split('*')
    if len(dim1) != 2 or len(dim2) != 2:
        erro_print(traceback.extract_stack(),
                   'Matrix dim is wrong!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None, None
    if dim1[-1] == '1':
        dim1 = ['1', dim1[0]]
    if dim2[-1] == '1':
        dim2 = ['1', dim2[0]]
    if dim1[0] != dim2[0] or dim1[-1] != dim2[-1]:
        erro_print(traceback.extract_stack(),
                   'Matrix dim is wrong!({}:{},{}:{})'.format(mat1, str(dim1), mat2, str(dim2)))
        return None, None
    return int(dim1[0]), int(dim1[-1])


def get_matrix_transpose_dim(input_dim):
    return f"{input_dim.split('*')[1]}*{input_dim.split('*')[0]}"


def get_matrix_dim_and_type(data: Data, mat):
    # return: Tuple(dim, type).
    # e.g. return ("1*3", "float64")
    mat = mat.strip()
    _, mat = delete_bracket(mat)
    if mat in data.global_vars.keys():
        input_dim = data.global_vars[mat]["dim"]
        input_type = data.global_vars[mat]["type"]
    elif mat in data.tmp_dict.keys():
        input_dim = data.tmp_dict[mat]["dim"]
        input_type = data.tmp_dict[mat]["type"]
    elif mat in data.const_dict.keys():
        input_dim = data.const_dict[mat]["dim"]
        input_type = data.const_dict[mat]["type"]
    else:
        erro_print(traceback.extract_stack(), 'Matrix {} not found!'.format(mat))
        return None, None
    return input_dim, input_type


def get_list_from_latex_matrix(expr):
    pass


def process_matrix_transpose(data: Data, node: Node):
    mat = node.children[0].code.strip()
    input_dim, input_type = get_matrix_dim_and_type(data, mat)
    if input_dim == None and input_type == None:
        return
    if input_dim.split('*')[0] == '1' or input_dim.split('*')[-1] == '1':
        # already generate the assignment-statement
        if node.father.op == 'none':
            return
        # input_dim = input_dim.split('*')[-1] + '*' + input_dim.split('*')[0]
        # if mat in data.global_vars.keys():
        #     data.global_vars[mat]["dim"] = input_dim
        # elif mat in data.tmp_dict.keys():
        #     data.tmp_dict[mat]["dim"] = input_dim
        # node.code = mat
        # node.set_op('matrix')
        # return
    res, func = '', ''
    for f in data.transpose_dict.keys():
        if input_dim == data.transpose_dict[f]["input_dim"] and input_type == data.transpose_dict[f]["input_type"]:
            func = f
            break
    if func == '':
        # 这里兼容了一下
        if input_type in ['siint32', 'unint32']:
            func = list(data.transpose_dict.keys())[1]
        elif input_type in ['float64', 'float32']:
            func = list(data.transpose_dict.keys())[0]
        else:
            # 默认用float，不报错
            func = list(data.transpose_dict.keys())[0]
    if node.father.op == '=':
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': input_type, 'dim': get_matrix_transpose_dim(input_dim)}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': input_type, 'dim': input_dim}
            res = get_matrix_pointer(data, var_name)
            node.father.children[0].code = var_name
        if func in list(data.transpose_dict.keys())[0:2]:
            node.set_code('{}({}, {}, {}, {})'.format(func, res, get_matrix_pointer(data, node.children[0].code), \
                                                      input_dim.split('*')[0], input_dim.split('*')[1]))
        else:
            node.set_code('{}({}, {})'.format(func, res, get_matrix_pointer(data, node.children[0].code)))
        node.set_op('hook')
        return
    else:
        # find the first ancestor of node whose op is 'body'
        fa = node.father
        son = node
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': input_type, 'dim': get_matrix_transpose_dim(input_dim)}
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                pres = get_matrix_pointer(data, res)
                if func in list(data.transpose_dict.keys())[0:2]:
                    func_line.set_code(
                        '{}({}, {}, {}, {})'.format(func, pres, get_matrix_pointer(data, node.children[0].code), \
                                                    input_dim.split('*')[0], input_dim.split('*')[1]))
                else:
                    func_line.set_code('{}({}, {})'.format(func, pres, get_matrix_pointer(data, node.children[0].code)))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        # modify the code
        pos = expr_pow(node.expr)
        l1, r1, l2, r2 = find_scope_pow(node.expr, pos)
        node.set_code((node.expr[:l1 - 1] + ' ' + res + ' ' + node.expr[r2 + 1:]).strip())
        ## To Be Modified: now suppose node.code is res
        node.set_op('matrix')


def process_matrix_inversion(data: Data, node: Node):
    mat = node.children[0].code.strip()
    input_dim, input_type = get_matrix_dim_and_type(data, mat)
    if input_dim == None and input_type == None:
        return
    if input_dim.split('*')[0] != input_dim.split('*')[-1]:
        erro_print(traceback.extract_stack(), 'Can not be inversed'.format(mat))
    res, func = '', ''
    for f in data.inversion_dict.keys():
        if input_dim == data.inversion_dict[f]["input_dim"] and input_type == data.inversion_dict[f]["input_type"]:
            output_dim = data.inversion_dict[f]["output_dim"]
            output_type = data.inversion_dict[f]["output_type"]
            func = f
            break
    if func == '':
        f = list(data.inversion_dict.keys())[0]
        if input_type == data.inversion_dict[f]["input_type"]:
            func = f
            output_type = input_type
            output_dim = input_dim
    if node.father.op == '=':
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': input_type, 'dim': input_dim}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': input_type, 'dim': input_dim}
            res = get_matrix_pointer(data, var_name)
            node.father.children[0].code = var_name
        node.set_code('{}({}, {}, {})'.format(func, res, get_matrix_pointer(data, node.children[0].code),
                                              input_dim.split('*')[0]))
        node.set_op('hook')
        return
    else:
        # find the first ancestor of node whose op is 'body'
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': output_type, 'dim': output_dim}
        fa = node.father
        son = node
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                pres = get_matrix_pointer(data, res)
                func_line.set_code('{}({}, {}, {})'.format(func, pres, get_matrix_pointer(data, node.children[0].code),
                                                           input_dim.split('*')[0]))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        # modify the code
        pos = expr_pow(node.expr)
        l1, r1, l2, r2 = find_scope_pow(node.expr, pos)
        node.set_code(
            (node.expr[:l1 - 1] + ' ' + res + ' ' + node.expr[r2 + 1:]).strip())
        ## To Be Modified: now suppose node.code is res
        node.set_op('matrix')


def process_matrix_multiply(data: Data, node: Node, l2r: bool = False):
    c1, c2 = node.children[0], node.children[1]
    dim1, _ = get_matrix_dim_and_type(data, c1.code)
    if dim1 == None and _ == None:
        return
    dim2, _ = get_matrix_dim_and_type(data, c2.code)
    if dim2 == None and _ == None:
        return
    if dim1.split('*')[1] == '1' and '1' not in dim2.split('*'):
        dim1 = f"{dim1.split('*')[1]}*{dim1.split('*')[0]}"
    elif dim2.split('*')[0] == '1' and '1' not in dim1.split('*'):
        dim2 = f"{dim2.split('*')[1]}*{dim2.split('*')[0]}"
    res_dim = matrix_multiply_dims_match(dim1, dim2, c1.code.strip(), c2.code.strip())
    if res_dim == None:
        return
    func, res = '', ''
    out_dim, out_type = '', ''
    for f in data.multi_dict.keys():
        if data.multi_dict[f]['input_dim'] == res_dim:
            func = f
            out_dim = data.multi_dict[f]['output_dim']
            out_type = data.multi_dict[f]['output_type']
            break
    if func == '':
        f = list(data.multi_dict.keys())[0]
        func = f
        out_dim = '{}*{}'.format(res_dim.split('*')[0], res_dim.split('*')[-1])
        out_type = data.multi_dict[f]['output_type']
    if node.father.op == '=' and not l2r:
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': out_type, 'dim': out_dim}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': out_type, 'dim': out_dim}
            res = var_name
            node.father.children[0].code = var_name
        elif var_dicts[res]['dim'] != out_dim:
            erro_print(traceback.extract_stack(), f"Matrix dim not match! {res} = {c1.code}*{c2.code}")
        res = get_matrix_pointer(data, res)
        if func == list(data.multi_dict.keys())[0]:
            node.set_code('{}({}, {}, {}, {}, {}, {})'.format( \
                func, res, get_matrix_pointer(data, c1.code), get_matrix_pointer(data, c2.code), res_dim.split('*')[0],
                res_dim.split('*')[1], res_dim.split('*')[2]))
        else:
            node.set_code('{}({}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code),
                                                  get_matrix_pointer(data, c2.code)))
        node.set_op('hook')
        return
    else:
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': out_type, 'dim': out_dim}
        # find the first ancestor of node whose op is 'body'
        fa = node.father
        son = node
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                pres = get_matrix_pointer(data, res)
                if func == list(data.multi_dict.keys())[0]:
                    func_line.set_code('{}({}, {}, {}, {}, {}, {})'.format( \
                        func, pres, get_matrix_pointer(data, c1.code), get_matrix_pointer(data, c2.code), \
                        res_dim.split('*')[0], res_dim.split('*')[1], res_dim.split('*')[2]))
                else:
                    func_line.set_code('{}({}, {}, {})'.format(func, pres, get_matrix_pointer(data, c1.code),
                                                               get_matrix_pointer(data, c2.code)))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        node.set_code(res)
        ## To Be Modified: now suppose node.code is res
        node.set_op('matrix')


def process_matrix_multiScalar(data: Data, node: Node, l2r: bool = False):
    c1, c2 = node.children[0], node.children[1]
    if c2.op == 'matrix':
        c1, c2 = c2, c1
    dim1, _ = get_matrix_dim_and_type(data, c1.code)
    if dim1 == None and _ == None:
        return
    res_dim = matrix_multiScalar_dim(dim1, c1.code.strip())
    if res_dim == None:
        return
    func, res = '', ''
    if res_dim[0] == '1':
        func = 'VectorScalar'
    else:
        func = 'MatrixScalar'
    out_dim, out_type = '{}*{}'.format(res_dim[0], res_dim[1]), data.multiScalar_dict[func]['output_type']
    if node.father.op == '=' and not l2r:
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': out_type, 'dim': out_dim}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': out_type, 'dim': out_dim}
            res = var_name
            node.father.children[0].code = var_name
        res = get_matrix_pointer(data, res)
        if func == 'VectorScalar':
            node.set_code(
                '{}({}, {}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code), c2.code, res_dim[1]))
        else:
            node.set_code(
                '{}({}, {}, {}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code), c2.code, res_dim[0],
                                                res_dim[1]))
        node.set_op('hook')
        return
    else:
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': out_type, 'dim': out_dim}
        # find the first ancestor of node whose op is 'body'
        fa = node.father
        son = node
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                pres = get_matrix_pointer(data, res)
                if func == 'VectorScalar':
                    func_line.set_code(
                        '{}({}, {}, {}, {})'.format(func, pres, get_matrix_pointer(data, c1.code), c2.code, res_dim[1]))
                else:
                    func_line.set_code(
                        '{}({}, {}, {}, {}, {})'.format(func, pres, get_matrix_pointer(data, c1.code), c2.code,
                                                        res_dim[0], res_dim[1]))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        node.set_code(res)
        ## To Be Modified: now suppose node.code is res
        node.set_op('matrix')


def process_matrix_add_and_sub(op: str, data: Data, node: Node, flag: bool = False):
    c1, c2 = node.children[0], node.children[1]
    dim1, _ = get_matrix_dim_and_type(data, c1.code)
    if dim1 == None and _ == None:
        return
    dim2, _ = get_matrix_dim_and_type(data, c2.code)
    if dim2 == None and _ == None:
        return
    res_dim = matrix_add_sub_dims_match(dim1, dim2, c1.code.strip(), c2.code.strip())
    if res_dim == None:
        return
    func, res = '', ''
    ty = ''
    if op == 'add':
        for f in data.add_dict.keys():
            if data.add_dict[f]['dim'] == res_dim:
                func = f
                ty = data.add_dict[f]['type']
                break
        if func == '':
            func = list(data.add_dict.keys())[0]
            ty = data.add_dict[func]['type']
    elif op == 'sub':
        for f in data.sub_dict.keys():
            if data.sub_dict[f]['dim'] == res_dim:
                func = f
                ty = data.sub_dict[f]['type']
                break
        if func == '':
            func = list(data.sub_dict.keys())[0]
            ty = data.sub_dict[func]['type']
    if node.father.op == '=' and not flag:
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': ty, 'dim': res_dim}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': ty, 'dim': res_dim}
            res = var_name
            node.father.children[0].code = var_name
        res = get_matrix_pointer(data, res)
        if op == 'add':
            if func == list(data.add_dict.keys())[0]:
                node.set_code('{}({}, {}, {}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code), \
                                                              get_matrix_pointer(data, c2.code), res_dim.split('*')[0],
                                                              res_dim.split('*')[1]))
            else:
                node.set_code('{}({}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code),
                                                      get_matrix_pointer(data, c2.code)))
        elif op == 'sub':
            if func == list(data.sub_dict.keys())[0]:
                node.set_code('{}({}, {}, {}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code), \
                                                              get_matrix_pointer(data, c2.code), res_dim.split('*')[0],
                                                              res_dim.split('*')[1]))
            else:
                node.set_code('{}({}, {}, {})'.format(func, res, get_matrix_pointer(data, c1.code),
                                                      get_matrix_pointer(data, c2.code)))
        node.set_op('hook')
        return
    else:
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': ty, 'dim': res_dim}
        # find the first ancestor of node whose op is 'body'
        fa = node.father
        son = node
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                p_res = get_matrix_pointer(data, res)
                if op == 'add':
                    if func == list(data.add_dict.keys())[0]:
                        func_line.set_code(
                            '{}({}, {}, {}, {}, {})'.format(func, p_res, get_matrix_pointer(data, c1.code),
                                                            get_matrix_pointer(data, c2.code), \
                                                            res_dim.split('*')[0], res_dim.split('*')[1]))
                    else:
                        func_line.set_code('{}({}, {}, {})'.format(func, p_res, get_matrix_pointer(data, c1.code),
                                                                   get_matrix_pointer(data, c2.code)))
                elif op == 'sub':
                    if func == list(data.sub_dict.keys())[0]:
                        func_line.set_code(
                            '{}({}, {}, {}, {}, {})'.format(func, p_res, get_matrix_pointer(data, c1.code),
                                                            get_matrix_pointer(data, c2.code), \
                                                            res_dim.split('*')[0], res_dim.split('*')[1]))
                    else:
                        func_line.set_code('{}({}, {}, {})'.format(func, p_res, get_matrix_pointer(data, c1.code),
                                                                   get_matrix_pointer(data, c2.code)))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        node.set_code(res)
        ## To Be Modified: now suppose pos_node.code is res
        node.set_op('matrix')


def process_matrix_equal(data: Data, node: Node):
    c1, c2 = node.children[0], node.children[1]
    dim1, dim2 = '', ''
    row, col = -1, -1
    dim1, _ = get_matrix_dim_and_type(data, c1.code)
    if dim1 == None and _ == None:
        return
    dim2, _ = get_matrix_dim_and_type(data, c2.code)
    if dim2 == None and _ == None:
        return
    row, col = matrix_dims_match(dim1, dim2, c1.code.strip(), c2.code.strip())
    if row == None and col == None:
        return
    func = list(data.equal_dict.keys())[0]

    node.set_code('{}({}, {}, {}, {})'.format(func, \
                                              get_matrix_pointer(data, c1.code), get_matrix_pointer(data, c2.code),
                                              str(row), str(col)))


def process_vector_matrix(data: Data, node: Node, vars):
    # node's expr is a vector or matrix which consists of vectors(20240213规范第11条)

    if node.father.op == '=':
        vec = node.father.children[0].code
        node.father.set_op('none')
    elif node.father.op == 'pow' and node.father.children[-1].expr == 'T' and node.father.father.op == '=':
        vec = node.father.father.children[0].code
        node.father.father.set_op('none')
    else:
        vec = data.get_tmp_var()
    vec_type, vec_dim = '', '1*{}'.format(str(len(vars)))
    if node.expr.find(';') != -1 or node.expr.find(r'\n') != -1:
        vec_dim = '{}*1'.format(str(len(vars)))
    element_dim = '1*1'
    var = vars[0].split('[')[0]  # 数组元素取其名，例：a[0]
    if var in data.global_vars.keys():
        vec_type = data.global_vars[var]['type']
        element_dim = data.global_vars[var]['dim']
        if vars[0].find('[') != -1:
            element_dim = '1*1'
    elif var in data.tmp_dict.keys():
        vec_type = data.tmp_dict[var]['type']
        element_dim = data.tmp_dict[var]['dim']
        if vars[0].find('[') != -1:
            element_dim = '1*1'
    elif var in data.const_dict.keys():
        vec_type = data.const_dict[var]['type']
        element_dim = data.const_dict[var]['dim']
        if vars[0].find('[') != -1:
            element_dim = '1*1'
    else:
        t = node
        while vec_type == '' and t.children != []:
            for c in t.children:
                if c.code in data.global_vars.keys():
                    vec_type = data.global_vars[c.code]['type']
                    element_dim = data.global_vars[c.code]['dim']
                    break
                elif c.code in data.tmp_dict.keys():
                    vec_type = data.tmp_dict[c.code]['type']
                    element_dim = data.tmp_dict[c.code]['dim']
                    break
                elif c.code in data.const_dict.keys():
                    vec_type = data.const_dict[c.code]['type']
                    element_dim = data.const_dict[c.code]['dim']
                    break
            t = c
        if vec_type == '':
            vec_type = 'float64'
        # data.tmp_dict[vec] = {'type':vec_type, 'dim':vec_dim}
    var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
    is_matrix = False
    if element_dim != '1*1':
        # 应为行向量或列向量
        if '1' not in element_dim:
            erro_print(traceback.extract_stack(), f'不符合矩阵表示规范！{node.expr}')
        if element_dim.split('*')[0] == '1':
            vec_dim = f'{str(len(vars))}*{element_dim.split("*")[1]}'
        else:
            vec_dim = f'{element_dim.split("*")[0]}*{str(len(vars))}'
        is_matrix = True
    if vec not in var_dicts.keys():
        var_name = data.get_recommend_name(vec)
        data.tmp_dict[var_name] = {'type': vec_type, 'dim': vec_dim}
        if vec != var_name:
            data.tmp_dict[vec] = {'name': var_name, 'type': vec_type, 'dim': vec_dim}
        vec = var_name
    fa = node.father
    son = node
    while fa.op != 'body':
        son = fa
        fa = fa.father
    for i in range(len(fa.children)):
        if fa.children[i] == son:
            for j in range(len(vars)):
                if is_matrix:
                    if element_dim[0].split('*')[0] == '1':
                        col_dim = int(element_dim.split('*')[1])
                        for k in range(col_dim):
                            line = Node(op='=')
                            line.set_code(f'{vec}[{str(j)}][{str(k)}] = {vars[j]}[{str(k)}]')
                            fa.add_child(line, i + j * col_dim + k)
                            line.generate_code()
                    else:
                        line_dim = int(element_dim.split('*')[0])
                        for k in range(line_dim):
                            line = Node(op='=')
                            line.set_code(f'{vec}[{str(k)}][{str(j)}] = {vars[j]}[{str(k)}]')
                            fa.add_child(line, i + j * line_dim + k)
                            line.generate_code()
                else:
                    line = Node(op='=')
                    line.set_code('{}[{}] = {}'.format(vec, str(j), vars[j]))
                    fa.add_child(line, i + j)
                    line.generate_code()
            break
    node.set_code(vec)
    node.set_op('matrix')


def process_array(data: Data, node: Node, op):
    # 处理数组，暂时只支持任意维度变量取第一维，其它情况先不报错
    for var in data.global_vars.keys():
        if op.find(var) == 0:
            post = op[len(var):].strip()
            if post[0] == '[' and post[len(post) - 1] == ']':
                node.set_code(op)
                node_dim = node_dim = get_arrayElement_dim(var, op, data.global_vars[var]['dim'])
                data.global_vars[op] = {'type': data.global_vars[var]['type'], 'dim': node_dim}
                if node_dim != '1*1':
                    node.set_op('matrix')
                else:
                    node.set_op('var')
                return True
    for var in data.tmp_dict.keys():
        if op.find(var) == 0:
            post = op[len(var):].strip()
            if post[0] == '[' and post[len(post) - 1] == ']':
                node.set_code(op)
                node_dim = get_arrayElement_dim(var, op, data.tmp_dict[var]['dim'])
                # 放到运行时全局变量字典里，以便运算时快速获取信息
                data.global_vars[op] = {'type': data.tmp_dict[var]['type'], 'dim': node_dim}
                if node_dim != '1*1':
                    node.set_op('matrix')
                else:
                    node.set_op('var')
                return True
    for var in data.const_dict.keys():
        if op.find(var) == 0:
            post = op[len(var):].strip()
            if post[0] == '[' and post[len(post) - 1] == ']':
                node.set_code(op)
                node_dim = get_arrayElement_dim(var, op, data.const_dict[var]['dim'])
                data.const_dict[op] = {'type': data.const_dict[var]['type'], 'dim': node_dim}
                if node_dim != '1*1':
                    node.set_op('matrix')
                else:
                    node.set_op('var')
                return True
    return False

    # dim = ''
    # var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
    # for var in var_dicts.keys():
    #     if op.find(var) == 0:
    #         post = op[len(var):].strip()
    #         if post[0] == '[':
    #             if (dim := get_arrayElement_dim(var, op, var_dicts[var]['dim'])) != None:
    #                 if op not in var_dicts.keys():
    #                     data.global_vars[op] = {'type':var_dicts[var]['dim'], 'dim':dim}
    #                 node.set_code(op)
    #                 if dim == '1*1':
    #                     node.set_op('var')
    #                 else:
    #                     node.set_op('matrix')
    #                 return True
    # return False


def process_vector_norm_matrix_det(data: Data, node: Node):
    # node: |vec| or |matrix|

    vec = node.children[0].code
    dim, ty = get_matrix_dim_and_type(data, vec)
    dims = dim.split('*')
    if dims[0] == '1' and dims[1] == '1':
        erro_print(traceback.extract_stack(), 'Not a vector!({})'.format(vec))
    # TO BE MODIFIED. Only two function supported now
    func, N = '', ''
    if '1' in dims:
        if ty == 'float32':
            func = 'VectorNormNF'
        else:
            func = 'VectorNormN'
        if dims[0] == '1':
            N = dims[1]
        else:
            N = dims[0]
        line = f'{func}({get_matrix_pointer(data, vec)}, {N})'
        node.set_code(line)
        node.set_op('func')
    else:
        if not (dims[0] == '3' and dims[1] == '3'):
            erro_print(traceback.extract_stack(), 'Not a 3*3 matrix!({})'.format(vec))
        if ty == 'float64':
            func = 'MatrixDet33'
        elif ty == 'siint32':
            func = 'MatrixDet33Int'
        else:
            erro_print(traceback.extract_stack(),
                       'Matrix {} type {} not supported for calculating determinant!'.format(vec, ty))
        line = f'{func}({get_matrix_pointer(data, vec)})'
        node.set_code(line)
        node.set_op('func')


def process_vector_cross(data: Data, node: Node, l2r: bool = False):
    c1, c2 = node.children[0], node.children[1]
    dim1, ty1 = get_matrix_dim_and_type(data, c1.code)
    if dim1 == None and ty1 == None:
        return
    dim2, ty2 = get_matrix_dim_and_type(data, c2.code)
    if dim2 == None and ty2 == None:
        return
    res_dim = matrix_add_sub_dims_match(dim1, dim2, c1.code.strip(), c2.code.strip())
    if res_dim == None:
        return
    func, res = '', ''

    if ty1 == 'float64':
        func = 'VectorCross3'
    elif ty1 == 'float32':
        func = 'VectorCross3F'
    if node.father.op == '=' and not l2r:
        res = node.father.children[0].code
        var_dicts = {**data.global_vars, **data.tmp_dict, **data.const_dict}
        if res not in var_dicts.keys():
            var_name = data.get_recommend_name(res)
            data.tmp_dict[var_name] = {'type': ty1, 'dim': res_dim}
            if res != var_name:
                data.tmp_dict[res] = {'name': var_name, 'type': ty1, 'dim': res_dim}
            res = var_name
            node.father.children[0].code = var_name
        res = get_matrix_pointer(data, res)
        node.set_code('{}({}, {}, {})'.format(func, res, c1.code, c2.code))
        node.set_op('hook')
        return
    else:
        res = data.get_tmp_var()
        data.tmp_dict[res] = {'type': ty1, 'dim': res_dim}
        # find the first ancestor of node whose op is 'body'
        fa = node.father
        son = node
        while fa.op != 'body':
            son = fa
            fa = fa.father
        for i in range(len(fa.children)):
            if fa.children[i] == son:
                func_line = Node(op='func')
                func_line.set_code('{}({}, {}, {})'.format(func, get_matrix_pointer(data, res), c1.code, c2.code))
                fa.add_child(func_line, i)
                func_line.generate_code()
                break
        node.set_code(res)
        ## To Be Modified: now suppose pos_node.code is res
        node.set_op('matrix')
